/*
 * Copyright(C) 2020. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <limits>
#include <omp.h>

#include <ascenddaemon/impl/TopkOp.h>
#include <ascenddaemon/utils/Limits.h>

namespace ascend {
TopkOp::TopkOp() {}

TopkOp::~TopkOp() {}

void TopkOp::reorder(AscendTensor<float16_t, DIMS_2> &topkDistance, AscendTensor<uint32_t, DIMS_2> &topkIndices)
{
    int n = topkIndices.getSize(0);

#pragma omp parallel for if (n > 1) num_threads(n < 6 ? n : 6)
    for (int i = 0; i < n; i++) {
        auto dist = topkDistance[i].view();
        auto id = topkIndices[i].view();
        reorder(dist, id);
    }
}

void TopkOp::reorder(AscendTensor<float16_t, 1> &topkDistance, AscendTensor<uint32_t, 1> &topkIndices)
{
    float16_t *dist = topkDistance.data();
    uint32_t *ids = topkIndices.data();
    size_t k = (size_t)topkIndices.getSize(0);

    size_t i, j;

    for (i = 0, j = 0; i < k; i++) {
        /* top element should be put at the end of the list */
        float16_t val = dist[0];
        uint32_t id = ids[0];

        maxValuePop(k - i, dist, ids);
        dist[k - j - 1] = val;
        ids[k - j - 1] = id;
        if (id != std::numeric_limits<uint32_t>::max()) {
            j++;
        }
    }

    auto err = memmove_s(dist, k * sizeof(*dist), dist + k - j, j * sizeof(*dist));
    ASCEND_THROW_IF_NOT_FMT(err == EOK, "An error occured when memmove in TopK reorder, %d!\n", err);
    err = memmove_s(ids, k * sizeof(*ids), ids + k - j, j * sizeof(*ids));
    ASCEND_THROW_IF_NOT_FMT(err == EOK, "An error occured when memmove in TopK reorder, %d!\n", err);

    for (; j < k; j++) {
        dist[j] = Limits<float16_t>::getMax();
        ids[j] = std::numeric_limits<uint32_t>::max();
    }
}

bool TopkOp::exec(AscendTensor<float16_t, DIMS_2> &distance,
                  AscendTensor<uint32_t, DIMS_2> &indices,
                  AscendTensor<float16_t, DIMS_2> &topkDistance,
                  AscendTensor<uint32_t, DIMS_2> &topkIndices,
                  uint32_t indexOffset)
{
    if (!checkParams(distance, indices, topkDistance, topkIndices)) {
        return false;
    }

    float16_t *disArray = distance.data();
    uint32_t *idArray = indices.data();
    float16_t *disResult = topkDistance.data();
    uint32_t *idResult = topkIndices.data();
    const int batch = distance.getSize(0);
    const int numDist = distance.getSize(1);
    const int kVal = topkDistance.getSize(1);
    bool idFlag = idArray ? true : false;

#pragma omp parallel for if (batch > 1) num_threads(batch < 6 ? batch : 6)
    for (int i = 0; i < batch; i++) {
        int batchOffset = numDist * i;
        int heapOffset = kVal * i;
        for (int j = 0; j < numDist; j++) {
            if (*(disArray + batchOffset + j) >= *(disResult + heapOffset)) {
                continue;
            }
            int index = idFlag ? *(idArray + batchOffset + j) : (indexOffset + j);
            minValuePush(kVal, disResult + heapOffset, idResult + heapOffset, *(disArray + batchOffset + j), index);
        }
    }

    return true;
}

bool TopkOp::exec(AscendTensor<float16_t, DIMS_1> &distance,
                  AscendTensor<uint32_t, DIMS_1> &indices,
                  AscendTensor<float16_t, DIMS_1> &topkDistance,
                  AscendTensor<uint32_t, DIMS_1> &topkIndices,
                  uint32_t realSize,
                  uint32_t indexOffset)
{
    if (!checkParams(distance, indices, topkDistance, topkIndices)) {
        return false;
    }

    if (distance.getSize(0) < (int)realSize) {
        return false;
    }

    float16_t *disArray = distance.data();
    uint32_t *idArray = indices.data();
    float16_t *disResult = topkDistance.data();
    uint32_t *idResult = topkIndices.data();
    const int kVal = topkDistance.getSize(0);
    bool idFlag = idArray ? true : false;

    for (uint32_t j = 0; j < realSize; j++) {
        if (*(disArray + j) >= *(disResult)) {
            continue;
        }
        int index = idFlag ? *(idArray + j) : (indexOffset + j);
        minValuePush(kVal, disResult, idResult, *(disArray + j), index);
    }

    return true;
}

bool TopkOp::exec(AscendTensor<float16_t, 1> &distance,
                  AscendTensor<uint32_t, 1> &indices,
                  AscendTensor<float16_t, 1> &topkDistance,
                  AscendTensor<uint32_t, 1> &topkIndices)
{
    if (!checkParams(distance, indices, topkDistance, topkIndices)) {
        return false;
    }

    float16_t *disArray = distance.data();
    uint32_t *idArray = indices.data();
    float16_t *disResult = topkDistance.data();
    uint32_t *idResult = topkIndices.data();
    const int numDist = distance.getSize(0);
    const int kVal = topkDistance.getSize(0);
    bool idFlag = idArray ? true : false;

    for (int j = 0; j < numDist; j++) {
        if (*(disArray + j) >= *(disResult)) {
            continue;
        }
        int index = idFlag ? *(idArray + j) : j;
        minValuePush(kVal, disResult, idResult, *(disArray + j), index);
    }

    return true;
}

void TopkOp::maxValuePop(size_t k, float16_t *heapDist, uint32_t *heapId)
{
    heapDist--; /* Use 1-based indexing for easier node->child translation */
    heapId--;
    float16_t val = heapDist[k];
    size_t i = 1;
    size_t i1;
    size_t i2;
    while (1) {
        i1 = i << 1;
        i2 = i1 + 1;
        if (i1 > k) {
            break;
        }

        if (i2 == k + 1 || heapDist[i1] > heapDist[i2]) {
            if (val > heapDist[i1]) {
                break;
            }
            heapDist[i] = heapDist[i1];
            heapId[i] = heapId[i1];
            i = i1;
        } else {
            if (val > heapDist[i2]) {
                break;
            }
            heapDist[i] = heapDist[i2];
            heapId[i] = heapId[i2];
            i = i2;
        }
    }
    heapDist[i] = heapDist[k];
    heapId[i] = heapId[k];
}

void TopkOp::minValuePush(size_t size, float16_t *heapDist, uint32_t *heapId, float16_t pushDist, uint32_t pushId)
{
    size_t i = 1;
    size_t leftChild, rightChild;
    heapDist--;
    heapId--;
    while (1) {
        leftChild = i << 1;
        rightChild = leftChild + 1;
        if (leftChild > size) {
            break;
        }

        if (rightChild == (size + 1) || *(heapDist + leftChild) > *(heapDist + rightChild)) {
            if (pushDist > *(heapDist + leftChild)) {
                break;
            }

            *(heapDist + i) = *(heapDist + leftChild);
            *(heapId + i) = *(heapId + leftChild);
            i = leftChild;
        } else {
            if (pushDist > *(heapDist + rightChild)) {
                break;
            }

            *(heapDist + i) = *(heapDist + rightChild);
            *(heapId + i) = *(heapId + rightChild);
            i = rightChild;
        }
    }

    *(heapDist + i) = pushDist;
    *(heapId + i) = pushId;
}

bool TopkOp::checkParams(AscendTensor<float16_t, 1> &distance,
                         AscendTensor<uint32_t, 1> &indices,
                         AscendTensor<float16_t, 1> &topkDistance,
                         AscendTensor<uint32_t, 1> &topkIndices)
{
    if (distance.data() == nullptr || topkDistance.data() == nullptr || topkIndices.data() == nullptr) {
        return false;
    }

    int outDistSize = topkDistance.getSize(0);
    int outIdSize = topkIndices.getSize(0);
    if (outIdSize != outDistSize) {
        return false;
    }

    return true;
}

bool TopkOp::checkParams(AscendTensor<float16_t, DIMS_2> &distance,
                         AscendTensor<uint32_t, DIMS_2> &indices,
                         AscendTensor<float16_t, DIMS_2> &topkDistance,
                         AscendTensor<uint32_t, DIMS_2> &topkIndices)
{
    if (distance.data() == nullptr || topkDistance.data() == nullptr || topkIndices.data() == nullptr) {
        return false;
    }

    int inDistSize0 = distance.getSize(0);
    int outDistSize0 = topkDistance.getSize(0);
    int outDistSize1 = topkDistance.getSize(1);
    int outIdSize0 = topkIndices.getSize(0);
    int outIdSize1 = topkIndices.getSize(1);
    if (inDistSize0 != outDistSize0 || inDistSize0 != outIdSize0 || outIdSize1 != outDistSize1) {
        return false;
    }

    return true;
}
} // namespace ascend
